import json

import torch
import torchvision.transforms

from backbone.darknet53 import YoloBody
from backbone.d import YOLOModel
from utils.dataset import COCO
from utils.object_detection import bbox_position_cast, non_max_suppress
from utils.train import init_weights
from framework.yolov3 import YOLOLoss, YOLOInfer
import torch.utils.data

if __name__ == '__main__':
    device = torch.device('cpu')
    transform = torchvision.transforms.Compose([
        torchvision.transforms.ToTensor(),
        torchvision.transforms.Resize((416, 416)),
    ])

    train_dataset = COCO('./out/data/datasets/banana-detection/train', 'images', 'annotations.json',
                         transform=transform)

    max_anchors = torch.tensor([[116, 90],
                                [156, 198],
                                [373, 326]], device=device)
    medium_anchors = torch.tensor([[30, 61],
                                   [62, 45],
                                   [59, 119],
                                   ], device=device)
    min_anchors = torch.tensor([[10, 13],
                                [16, 30],
                                [33, 23], ], device=device)

    yolo_loss = YOLOLoss([max_anchors, medium_anchors, min_anchors], c=1, strides=[13, 26, 52],
                         image_size=torch.Size([416, 416]), device=device)

    model = YoloBody([3, 3, 3], 1)

    epoch = 30

    model.apply(init_weights)
    model.to(device)


    def collate_fn(batch):
        images = []
        labels = []
        for image, label in batch:
            images.append(image)
            labels.append(label)

        return torch.stack(images), labels


    model.train()

    optimizer = torch.optim.Adam(model.parameters(), 1e-2, weight_decay=5e-4)

    lr_scheduler = torch.optim.lr_scheduler.StepLR(optimizer, step_size=1, gamma=0.94)

    for e in range(epoch):
        i = 0
        L = 0.

        for x, y in torch.utils.data.DataLoader(train_dataset, batch_size=30, shuffle=True, collate_fn=collate_fn):
            x = x.to(device)
            for i in range(len(y)):
                y[i] = y[i].to(device)

            optimizer.zero_grad()
            y0, y1, y2 = model(x)
            y_p = [y0, y1, y2]
            l = yolo_loss(y_p, y)
            l.backward()
            optimizer.step()
            L += l.item()
            i += 1
            break
        print(L / i)
        lr_scheduler.step()

    torch.save(model.state_dict(), './train_e30.pth')
